# -*- coding: utf-8 -*-

"""
@File    : 07_文件转换之IMG转TIF-批量.py
@Author  : fungis@163.com
@notice  :
"""

import os

from osgeo import gdal, gdalconst


filter_list = ['.IMG', '.img']  # 设置过滤后的文件类型 当然可以设置多个类型
def find_filter_file(dirname):
    result = []  # 所有的文件
    for maindir, subdir, file_name_list in os.walk(dirname):
        # print("提示：当前主目录下的所有目录", "2:", subdir)  # 当前主目录下的所有目录
        # print("3:", file_name_list)  # 当前主目录下的所有文件

        for filename in file_name_list:
            apath = os.path.join(maindir, filename)  # 合并成一个完整路径
            ext = os.path.splitext(apath)[1]  # 获取文件后缀 [0]获取的是除了文件名以外的内容
            # 设置过滤器，过滤不符合格式的文件
            if ext in filter_list:
                result.append(apath)
    return result

def gis_img_to_tif(in_data, out_path, dtype):
    """
    保存tif文件
    :param in_data: 原始数据路径
    :param out_path: 保存tif文件路径
    :param dtype: gdal数据类型, 默认自动根据输入识别
    :return:
    """
    if os.path.exists(in_data) is False:
        raise Exception('[Errno 2] 该文件不存在: \'' + in_data + '\'')

    dataset = gdal.Open(in_data, gdalconst.GA_ReadOnly)  # 读取遥感数据信息
    band_count = dataset.RasterCount  # 波段数
    if (band_count > 0):
        img_to_tif_write(dataset, band_count, dtype, out_path)
    else:
        raise Exception('[Errno 4] 该文件波段异常: \'' + in_data + '\'')

def img_to_tif_write(dataset, band_count, dtype, out_path):
    geo_transform = dataset.GetGeoTransform()  # 获取地理空间信息
    columns = dataset.RasterXSize
    row = dataset.RasterYSize
    dim_z = band_count  # 波段数

    # 创建.tif文件
    driver = gdal.GetDriverByName('GTiff')  # 创建驱动
    # 创建文件
    if dtype == "float" or dtype == "float32":
        dst_ds = driver.Create(out_path, columns, row, dim_z, gdal.GDT_Float32)
    elif dtype == "float64":
        dst_ds = driver.Create(out_path, columns, row, dim_z, gdal.GDT_Float64)
    elif dtype == "int":
        dst_ds = driver.Create(out_path, columns, row, dim_z, gdal.GDT_UInt16)
    else:
        raise Exception('[Errno 4] 未设置值类型或设置字符串存在错误！')
    dst_ds.SetGeoTransform(geo_transform)  # 设置几何信息
    dst_ds.SetProjection(dataset.GetProjection())  # 设置投影
    for i in range(1, (int(band_count) + 1)):
        band = dataset.GetRasterBand(i)
        cell_map = band.ReadAsArray()
        dst_ds.GetRasterBand(i).WriteArray(cell_map)  # 写入像元信息
        dst_ds.FlushCache()  # 写入硬盘
    del dst_ds  # 关闭tif文件

if __name__ == '__main__':

    input_file_path = r'./data-use/img/'  # 输入路径
    output_file_path = r'./results/img/'  # 输出路径
    if not os.path.exists(output_file_path):
        # os.makedirs可以创建路径上多个
        os.makedirs(output_file_path)
    file_list = find_filter_file(input_file_path)
    if len(file_list) > 0:
        for raster_path in file_list:
            # 输出文件路径
            name = os.path.basename(raster_path).split('.')[0]
            result_file = output_file_path + name + '.tif'  # 输出文件位置
            del name  # 删除无效变量
            # 执行转换函数
            gis_img_to_tif(raster_path, result_file, 'float')
        print('------处理完成------')
    else:
        raise Exception('[Errno 1] 该路径下无{}类型文件'.format(filter_list))